/*
 *  Copyright 2008-2009 NVIDIA Corporation
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */


#include <algorithm>
#include <vector>
#include <stdio.h>
#include <math.h>

#include <thrust/adjacent_difference.h>

#include <cusp/graph/detail/host/metis.h>
#include <cusp/graph/detail/host/patoh.h>

namespace cusp
{
namespace graph
{
namespace detail
{
namespace host
{

template <class MatrixType, class ArrayType1, class ArrayType2>
size_t partition(const MatrixType& A,
		 size_t m, 
                 size_t num_parts,
                 const ArrayType1& Cwghts,
		 ArrayType2& part,
                 cusp::graph::blk)
{
    typedef typename MatrixType::index_type IndexType;
    typedef typename MatrixType::value_type ValueType;

    part.resize(m, 0);
    int n = ceil(double(m)/double(num_parts));

    for( int i = 0; i < int(m); i++ )
	part[i] = floor(i/n);

    return 0;
}

template <class MatrixType, class ArrayType1, class ArrayType2>
size_t partition(const MatrixType& A,
		 size_t m, 
                 size_t num_parts,
                 const ArrayType1& Cwghts,
		 ArrayType2& part,
                 cusp::graph::wbk)
{
    typedef typename MatrixType::index_type IndexType;
    typedef typename MatrixType::value_type ValueType;

    part.resize(m, 0);
    int n = ceil(double(m)/double(num_parts));

    ValueType sum = 0;
    for( int i = 0; i < int(m); i++ )
	sum += Cwghts[i];

    int avgw = sum / num_parts;
    int curw = 0;
    int part_num = 0;
    part.resize(m, 0);

    for( int i = 0; i < int(m); i++ ){
	part[i] = part_num;
	curw += Cwghts[i];
	if( curw > avgw )
	{
		part_num += 1;
		curw = 0;
	}
    }

    return 0;
}

template <typename PartitionerType, class MatrixType, class ArrayType>
size_t partition(const MatrixType& graph,
                 size_t num_parts,
                 ArrayType& part,
                 bool Kway,
		 cusp::graph::rowwise)
{
    typedef typename MatrixType::index_type IndexType;
    typedef typename MatrixType::value_type ValueType;

    //bool is_symmetric = cusp::graph::is_structurally_symmetric(graph);

    size_t M = graph.num_rows;

    cusp::csr_matrix<IndexType,ValueType,cusp::host_memory> A(graph);
    
    cusp::array1d<ValueType,cusp::host_memory> Cwghts(M);
    thrust::adjacent_difference( A.row_offsets.begin()+1, A.row_offsets.end(), Cwghts.begin() );
    
    return partition( A, M, num_parts, Cwghts, part, PartitionerType() );  
}

template <typename PartitionerType, class MatrixType, class ArrayType>
size_t partition(const MatrixType& graph,
                 size_t num_parts,
                 ArrayType& part,
                 bool Kway,
		 cusp::graph::colwise)
{
    typedef typename MatrixType::index_type IndexType;
    typedef typename MatrixType::value_type ValueType;

    //bool is_symmetric = cusp::graph::is_structurally_symmetric(graph);

    size_t N = graph.num_cols;

    cusp::csr_matrix<IndexType,ValueType,cusp::host_memory> A(graph);
    cusp::csr_matrix<IndexType,ValueType,cusp::host_memory> At;
    cusp::transpose(A,At);
    
    cusp::array1d<ValueType,cusp::host_memory> Cwghts(N);
    thrust::adjacent_difference( At.row_offsets.begin()+1, At.row_offsets.end(), Cwghts.begin() );
    
    return partition( At, N, num_parts, Cwghts, part, PartitionerType() );  
}

} // end namespace host
} // end namespace detail
} // end namespace graph
} // end namespace cusp

